package com.keba.lamp.ai.strategy.impl.doubao;

import cn.hutool.core.convert.Convert;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONUtil;
import com.baidubce.util.JsonUtils;
import com.google.common.collect.Lists;
import com.keba.lamp.ai.strategy.AbstractApiCallStrategy;
import com.keba.lamp.ai.vo.query.AiCallModelQuery;
import com.keba.lamp.ai.vo.query.AiCallModelRequestQuery;
import com.keba.lamp.ai.vo.result.AlibabaAiAudioResultVO;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionContentPart;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import okhttp3.ConnectionPool;
import okhttp3.Dispatcher;
import org.springframework.core.env.Environment;
import org.springframework.stereotype.Component;

import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;


@Slf4j
@Component("doubao")
public class DoubaoApiCallStrategy extends AbstractApiCallStrategy {

    @Resource
    private Environment environment;
    private ArkService arkService = null;

    private int maxTokens = 80;


    @PostConstruct
    public void init() {
        maxTokens = Convert.toInt(environment.getProperty("lamp.ai.doubao.max_images", "80"));
        arkService = ArkService.builder().dispatcher(new Dispatcher()).connectionPool(new ConnectionPool(5, 1, TimeUnit.SECONDS)).apiKey(environment.getProperty("lamp.ai.doubao.appKey")).build();
    }

    @Override
    public Object call(AiCallModelQuery aiCallModel) {
        List<ChatCompletionContentPart> parts = new ArrayList<>();
        aiCallModel.getCompletionContents().stream().filter(o -> StrUtil.isNotBlank(o.getType()) && (o.getType().equals("text") || o.getType().equals("image_url"))).forEach(o -> {
            ChatCompletionContentPart completionContentPart = ChatCompletionContentPart.builder().type(o.getType()).text(o.getContent()).content(o.getContent()).build();
            if (StrUtil.isNotBlank(o.getFileUrl()) && "image_url".equals(o.getType())) {
                completionContentPart.setImageUrl(new ChatCompletionContentPart.ChatCompletionContentPartImageURL(o.getFileUrl()));
            }
            parts.add(completionContentPart);
        });

        List<String> audioUrls = aiCallModel.getCompletionContents().stream().filter(o -> StrUtil.isNotBlank(o.getType()) && o.getType().equals("audio_url")).map(AiCallModelRequestQuery.CompletionContent::getFileUrl).collect(Collectors.toList());
        List<AlibabaAiAudioResultVO> audioResults = callAudio(audioUrls);
        audioResults.forEach(o -> {
            parts.add(ChatCompletionContentPart.builder().type("text").text(o.getTranscripts().get(0).getText()).build());
        });
        List<ChatMessage> messages = new ArrayList<>();
        if (StrUtil.isNotBlank(aiCallModel.getSystemPrompt())) {
            messages.add(ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(aiCallModel.getSystemPrompt()).build());
        }
        if (StrUtil.isNotBlank(aiCallModel.getUserPrompt())) {
            messages.add(ChatMessage.builder().role(ChatMessageRole.USER).content(aiCallModel.getUserPrompt()).build());
        }
        if (parts.size() <= maxTokens) {
            buildAiResponse(aiCallModel, parts, messages);
            return aiCallModel.getAiResponse();
        }

        return call(multiCall(parts, aiCallModel));
    }

    private void buildAiResponse(AiCallModelQuery aiCallModel, List<ChatCompletionContentPart> parts, List<ChatMessage> messages) {
        messages.add(ChatMessage.builder().role(ChatMessageRole.USER).multiContent(parts).build());
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder().model(aiCallModel.getModelId()).messages(messages).maxTokens(32288).build();
        log.info("请求豆包ai:{}", JsonUtils.toJsonString(chatCompletionRequest));
        arkService.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> {
            aiCallModel.setAiResponse(choice.getMessage().getContent());
            log.debug("DouBaoCall Response: >>> " + JsonUtils.toJsonString(choice));
        });
    }

    private AiCallModelQuery multiCall(List<ChatCompletionContentPart> parts, AiCallModelQuery aiCallModel) {
        List<ChatMessage> messages = new ArrayList<>();
        AiCallModelQuery aiCallModelQuery = AiCallModelQuery.builder().modelId(aiCallModel.getModelId()).systemPrompt(aiCallModel.getSystemPrompt()).userPrompt(aiCallModel.getUserPrompt()).build();
        List<AiCallModelRequestQuery.CompletionContent> completionContents = new ArrayList<>();
        List<List<ChatCompletionContentPart>> list = Lists.partition(parts, maxTokens);
        if (StrUtil.isNotBlank(aiCallModel.getSystemPrompt())) {
            messages.add(ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(aiCallModel.getSystemPrompt()).build());
        }
        messages.add(ChatMessage.builder().role(ChatMessageRole.USER).content("根据系统模版去掉多余的### 任务，编写法律相关案件，例如起诉状、答辩状等").build());
        list.forEach(o -> {
            messages.add(ChatMessage.builder().role(ChatMessageRole.USER).multiContent(o).build());
            ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder().model(aiCallModel.getModelId()).messages(messages).build();
            log.info("请求豆包ai:{}", JsonUtils.toJsonString(chatCompletionRequest));
            arkService.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> {
                completionContents.add(AiCallModelRequestQuery.CompletionContent.builder().type("text").content(choice.getMessage().getContent().toString()).build());
                log.debug("DouBaoCall Response: >>> " + JsonUtils.toJsonString(choice));
            });
        });

        aiCallModelQuery.setCompletionContents(completionContents);
        return aiCallModelQuery;
    }


}
